"""
测试模型架构模块
"""
import os
import sys
import unittest
import numpy as np
import torch
import torch.nn as nn

# 添加项目根目录到路径
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', '..')))

from sichuanmajiang.model.model import (
    MahjongModel, PolicyNetwork, ValueNetwork, ResidualBlock,
    create_mahjong_model, save_model, load_model, get_model_device
)
from sichuanmajiang.tests.test_utils import (
    get_test_device, create_test_features, create_test_directory, remove_test_directory
)


class TestResidualBlock(unittest.TestCase):
    """
    测试ResidualBlock类
    """
    def setUp(self):
        """
        测试前的设置
        """
        self.device = get_test_device()
        self.hidden_dim = 128
        self.res_block = ResidualBlock(self.hidden_dim).to(self.device)
    
    def test_forward_pass(self):
        """
        测试前向传播
        """
        # 创建输入
        batch_size = 16
        x = torch.randn(batch_size, self.hidden_dim, device=self.device)
        
        # 前向传播
        output = self.res_block(x)
        
        # 检查输出形状
        self.assertEqual(output.shape, (batch_size, self.hidden_dim))
    
    def test_residual_connection(self):
        """
        测试残差连接
        """
        # 暂时禁用激活函数，以便更容易验证残差连接
        self.res_block.activation = nn.Identity()
        
        # 创建输入
        batch_size = 8
        x = torch.randn(batch_size, self.hidden_dim, device=self.device)
        
        # 前向传播
        output = self.res_block(x)
        
        # 由于禁用了激活函数和批归一化（在测试环境中可能未启用），输出不应完全等于输入
        # 但应该是相关的
        self.assertTrue(torch.allclose(output, self.res_block.conv1(x) + self.res_block.conv2(self.res_block.bn1(self.res_block.conv1(x)))))


class TestPolicyNetwork(unittest.TestCase):
    """
    测试PolicyNetwork类
    """
    def setUp(self):
        """
        测试前的设置
        """
        self.device = get_test_device()
        self.input_dim = 324  # 特征维度
        self.action_dim = 6   # 动作空间大小
        self.num_res_blocks = 3
        self.hidden_dim = 128
        
        self.policy_net = PolicyNetwork(
            input_dim=self.input_dim,
            action_dim=self.action_dim,
            num_res_blocks=self.num_res_blocks,
            hidden_dim=self.hidden_dim
        ).to(self.device)
    
    def test_forward_pass(self):
        """
        测试前向传播
        """
        # 创建输入
        batch_size = 16
        x = torch.randn(batch_size, self.input_dim, device=self.device)
        
        # 前向传播
        output = self.policy_net(x)
        
        # 检查输出形状
        self.assertEqual(output.shape, (batch_size, self.action_dim))
    
    def test_softmax_output(self):
        """
        测试输出是否为有效的概率分布
        """
        # 创建输入
        batch_size = 16
        x = torch.randn(batch_size, self.input_dim, device=self.device)
        
        # 前向传播
        output = self.policy_net(x)
        
        # 检查是否所有元素都非负
        self.assertTrue(torch.all(output >= 0))
        
        # 检查每一行的和是否约等于1
        row_sums = torch.sum(output, dim=1)
        self.assertTrue(torch.allclose(row_sums, torch.ones(batch_size, device=self.device), atol=1e-5))


class TestValueNetwork(unittest.TestCase):
    """
    测试ValueNetwork类
    """
    def setUp(self):
        """
        测试前的设置
        """
        self.device = get_test_device()
        self.input_dim = 324  # 特征维度
        self.num_res_blocks = 3
        self.hidden_dim = 128
        
        self.value_net = ValueNetwork(
            input_dim=self.input_dim,
            num_res_blocks=self.num_res_blocks,
            hidden_dim=self.hidden_dim
        ).to(self.device)
    
    def test_forward_pass(self):
        """
        测试前向传播
        """
        # 创建输入
        batch_size = 16
        x = torch.randn(batch_size, self.input_dim, device=self.device)
        
        # 前向传播
        output = self.value_net(x)
        
        # 检查输出形状
        self.assertEqual(output.shape, (batch_size, 1))
    
    def test_output_range(self):
        """
        测试输出是否在合理范围内
        """
        # 创建输入
        batch_size = 16
        x = torch.randn(batch_size, self.input_dim, device=self.device)
        
        # 前向传播
        output = self.value_net(x).squeeze()
        
        # 检查值是否在合理范围内（例如-1到1之间）
        # 这取决于具体的实现，可能需要调整范围
        self.assertTrue(torch.all(output >= -1))
        self.assertTrue(torch.all(output <= 1))


class TestMahjongModel(unittest.TestCase):
    """
    测试MahjongModel类
    """
    def setUp(self):
        """
        测试前的设置
        """
        self.device = get_test_device()
        self.test_dir = os.path.join(os.path.dirname(__file__), 'test_models')
        create_test_directory(self.test_dir)
        
        self.input_dim = 324  # 特征维度
        self.action_dim = 6   # 动作空间大小
        self.num_res_blocks = 3
        self.hidden_dim = 128
        
        self.model = MahjongModel(
            input_dim=self.input_dim,
            action_dim=self.action_dim,
            num_res_blocks=self.num_res_blocks,
            hidden_dim=self.hidden_dim
        ).to(self.device)
    
    def tearDown(self):
        """
        测试后的清理
        """
        # 移除测试目录
        remove_test_directory(self.test_dir)
    
    def test_forward_pass(self):
        """
        测试前向传播
        """
        # 创建输入
        batch_size = 16
        x = torch.randn(batch_size, self.input_dim, device=self.device)
        
        # 前向传播
        policy_output, value_output = self.model(x)
        
        # 检查输出形状
        self.assertEqual(policy_output.shape, (batch_size, self.action_dim))
        self.assertEqual(value_output.shape, (batch_size, 1))
    
    def test_predict(self):
        """
        测试predict方法
        """
        # 创建输入
        x = torch.randn(1, self.input_dim, device=self.device)
        
        # 预测
        action_probs, value = self.model.predict(x)
        
        # 检查输出类型和形状
        self.assertIsInstance(action_probs, np.ndarray)
        self.assertIsInstance(value, float)
        self.assertEqual(action_probs.shape, (self.action_dim,))
    
    def test_save_and_load(self):
        """
        测试保存和加载模型
        """
        # 创建测试输入
        x = torch.randn(1, self.input_dim, device=self.device)
        
        # 保存前的输出
        original_policy, original_value = self.model(x)
        
        # 保存模型
        model_path = os.path.join(self.test_dir, 'test_model.pt')
        self.model.save(model_path)
        
        # 检查文件是否存在
        self.assertTrue(os.path.exists(model_path))
        
        # 加载模型
        loaded_model = MahjongModel.load(model_path)
        loaded_model.to(self.device)
        
        # 加载后的输出
        loaded_policy, loaded_value = loaded_model(x)
        
        # 检查输出是否相同
        self.assertTrue(torch.allclose(original_policy, loaded_policy))
        self.assertTrue(torch.allclose(original_value, loaded_value))
    
    def test_get_params(self):
        """
        测试获取模型参数
        """
        # 获取参数
        params = self.model.get_params()
        
        # 检查参数数量
        total_params = sum(p.numel() for p in self.model.parameters())
        self.assertEqual(params.numel(), total_params)
    
    def test_set_params(self):
        """
        测试设置模型参数
        """
        # 获取原始参数
        original_params = self.model.get_params().clone()
        
        # 创建新参数
        new_params = torch.randn_like(original_params)
        
        # 设置新参数
        self.model.set_params(new_params)
        
        # 检查参数是否已更新
        updated_params = self.model.get_params()
        self.assertTrue(torch.allclose(updated_params, new_params))
    
    def test_get_flattened_grads(self):
        """
        测试获取梯度
        """
        # 创建输入
        x = torch.randn(1, self.input_dim, device=self.device)
        
        # 前向传播
        policy, value = self.model(x)
        
        # 反向传播
        loss = policy.sum() + value.sum()
        loss.backward()
        
        # 获取梯度
        grads = self.model.get_flattened_grads()
        
        # 检查梯度是否存在且形状正确
        total_params = sum(p.numel() for p in self.model.parameters())
        self.assertEqual(grads.shape, (total_params,))
    
    def test_eval_mode(self):
        """
        测试评估模式
        """
        # 设置为评估模式
        self.model.eval()
        
        # 检查所有子模块是否都在评估模式
        for module in self.model.modules():
            if isinstance(module, (nn.BatchNorm1d, nn.Dropout)):
                self.assertFalse(module.training)
    
    def test_train_mode(self):
        """
        测试训练模式
        """
        # 设置为训练模式
        self.model.train()
        
        # 检查所有子模块是否都在训练模式
        for module in self.model.modules():
            if isinstance(module, (nn.BatchNorm1d, nn.Dropout)):
                self.assertTrue(module.training)


class TestHelperFunctions(unittest.TestCase):
    """
    测试辅助函数
    """
    def setUp(self):
        """
        测试前的设置
        """
        self.device = get_test_device()
        self.test_dir = os.path.join(os.path.dirname(__file__), 'test_models')
        create_test_directory(self.test_dir)
    
    def tearDown(self):
        """
        测试后的清理
        """
        # 移除测试目录
        remove_test_directory(self.test_dir)
    
    def test_create_mahjong_model(self):
        """
        测试create_mahjong_model函数
        """
        # 创建模型
        model = create_mahjong_model(
            input_dim=324,
            action_dim=6,
            num_res_blocks=3,
            hidden_dim=128
        )
        
        # 检查类型
        self.assertIsInstance(model, MahjongModel)
    
    def test_save_and_load_model_functions(self):
        """
        测试save_model和load_model函数
        """
        # 创建模型
        model = create_mahjong_model(
            input_dim=324,
            action_dim=6
        )
        model.to(self.device)
        
        # 创建测试输入
        x = torch.randn(1, 324, device=self.device)
        
        # 保存前的输出
        original_policy, original_value = model(x)
        
        # 保存模型
        model_path = os.path.join(self.test_dir, 'test_model_functions.pt')
        save_model(model, model_path)
        
        # 检查文件是否存在
        self.assertTrue(os.path.exists(model_path))
        
        # 加载模型
        loaded_model = load_model(model_path)
        loaded_model.to(self.device)
        
        # 加载后的输出
        loaded_policy, loaded_value = loaded_model(x)
        
        # 检查输出是否相同
        self.assertTrue(torch.allclose(original_policy, loaded_policy))
        self.assertTrue(torch.allclose(original_value, loaded_value))
    
    def test_get_model_device(self):
        """
        测试get_model_device函数
        """
        # 创建模型
        model = create_mahjong_model(
            input_dim=324,
            action_dim=6
        )
        
        # 默认设备应该是cpu
        self.assertEqual(get_model_device(model), torch.device('cpu'))
        
        # 移至其他设备
        if torch.cuda.is_available():
            model.to('cuda')
            self.assertEqual(get_model_device(model), torch.device('cuda'))
        
        # 移回cpu
        model.to('cpu')
        self.assertEqual(get_model_device(model), torch.device('cpu'))


if __name__ == '__main__':
    unittest.main()
